#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

void _create_sin_cos_table(int numangle, double min_theta, double theta_delta, std::vector<float>& sinTable, std::vector<float>& cosTable) {
    float angle = (float)min_theta;
    sinTable.resize(numangle);
    cosTable.resize(numangle);
    for (int i = 0; i < numangle; i++) {
        sinTable[i] = (float)std::sin(angle);
        cosTable[i] = (float)std::cos(angle);
        angle += (float)theta_delta;
    }
}

void _find_local_maximums(const cv::Mat& accum, int threshold, std::vector<int>& rs, std::vector<int>& cs) {
    rs.clear();
    cs.clear();
    for (int i = 1; i < accum.rows - 1; i++) {
        for (int j = 1; j < accum.cols - 1; j++) {
            int a = accum.ptr<int>(i)[j];
            if (a > threshold) {
                int left = accum.ptr<int>(i)[j-1];
                int right = accum.ptr<int>(i)[j+1];
                int top = accum.ptr<int>(i-1)[j];
                int bottom = accum.ptr<int>(i+1)[j];
                if (a > left && a >= right && a > top && a >= bottom) {
                    rs.push_back(i);
                    cs.push_back(j);
                }
            }
        }
    }
}

bool _compare_func(const cv::Vec3f& v1, const cv::Vec3f& v2) {
    return v1[2] > v2[2];
}

std::vector<cv::Vec3f> _my_hough_lines(const cv::Mat& binary, double rho, double theta, int threshold, double min_theta, double max_theta) {
    CV_Assert(binary.type() == CV_8UC1);
    
    int width = binary.cols;
    int height = binary.rows;
    int max_rho = (int)std::sqrt(width*width + height*height) + 1;
    int min_rho = -max_rho;
    
    int numangle = cvRound((max_theta - min_theta) / theta);
    int numrho = cvRound((max_rho - min_rho + 1) / rho);
    
    cv::Mat accum = cv::Mat::zeros(numangle+2, numrho+2, CV_32SC1);
    cv::Mat rhos = cv::Mat::zeros(numangle+2, numrho+2, CV_32FC1);
    cv::Mat angles = cv::Mat::zeros(numangle+2, numrho+2, CV_32FC1);
    
    std::vector<float> sinTable, cosTable;
    _create_sin_cos_table(numangle, min_theta, theta, sinTable, cosTable);
    
    for (int i = 0; i < height; i++) {
        for (int j = 0; j < width; j++) {
            if (binary.ptr<uchar>(i)[j] != 0) {
                for (int n = 0; n < numangle; n++) {
                    int r = cvRound((j*cosTable[n]+i*sinTable[n]) / rho);
                    r += (numrho - 1) / 2;
                    accum.ptr<int>(n+1)[r+1]++;
                    angles.ptr<float>(n+1)[r+1] = (float)min_theta + n*theta;
                    rhos.ptr<float>(n+1)[r+1] = cvRound(j*cosTable[n] + i*sinTable[n]);
                }
            }
        }
    }
    
    std::vector<int> rs, cs;
    _find_local_maximums(accum, threshold, rs, cs);
    
    std::vector<cv::Vec3f> lines;
    for (int i = 0; i < rs.size(); i++) {
        cv::Vec3f line;
        line[0] = rhos.ptr<float>(rs[i])[cs[i]];
        line[1] = angles.ptr<float>(rs[i])[cs[i]];
        line[2] = (float)accum.ptr<int>(rs[i])[cs[i]];
        lines.push_back(line);
    }
    
    std::sort(lines.begin(), lines.end(), _compare_func);
    
    return lines;
}

bool _is_result_same(const std::vector<cv::Vec3f>& v1, const std::vector<cv::Vec3f>& v2) {
    if (v1.size() != v2.size()) {
        return false;
    }
    if (v1.size() == 0) {
        return false;
    }
    for (int i = 0; i < v1.size(); i++) {
        if (std::abs(v1[i][0]-v2[i][0]) > FLT_EPSILON || std::abs(v1[i][1]-v2[i][1]) > FLT_EPSILON  || std::abs(v1[i][2]-v2[i][2]) > FLT_EPSILON) {
            return false;
        }
    }
    return true;
}

int main(int argc, char **argv) {
    cv::Mat src = cv::imread("/home/xlll/Downloads/opencv/samples/data/building.jpg", cv::IMREAD_GRAYSCALE);
    if (src.empty()) {
        std::cout << "failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    cv::Mat binary;
    cv::Canny(src, binary, 50, 230, 3, true);
    
    std::vector<cv::Vec3f> lines;
    double rho = 1;
    double theta = CV_PI / 180;
    double threshold = 200;
    double srn = 0, stn = 0;
    double min_theta = 0, max_theta = CV_PI;
    cv::HoughLines(binary, lines, rho, theta, threshold, srn, stn, min_theta, max_theta);
    std::cout << "number of lines = " << lines.size() << std::endl;
    cv::Mat dst(src.size(), CV_8UC1);
    for (int i = 0; i < lines.size(); i++) {
        double rho_i = lines[i][0], theta_i = lines[i][1], votes_i = lines[i][2];
        cv::Point pt1, pt2;
        double a = std::cos(theta_i), b = std::sin(theta_i);
        double x0 = a*rho_i, y0 = b*rho_i;
        pt1.x = std::round(x0 + 1000*(-b));
        pt1.y = std::round(y0 + 1000*a);
        pt2.x = std::round(x0 - 1000*(-b));
        pt2.y = std::round(y0 - 1000*a);
        cv::line(dst, pt1, pt2, cv::Scalar(255), 3, cv::LINE_AA);
    }
    std::cout << "================================" << std::endl;
    
    std::vector<cv::Vec3f> lines2;
    lines2 = _my_hough_lines(binary, rho, theta, threshold, min_theta, max_theta);
    std::cout << "number of lines2 = " << lines2.size() << std::endl;
    cv::Mat dst2(src.size(), CV_8UC1);
    for (int i = 0; i < lines2.size(); i++) {
        double rho_i = lines2[i][0], theta_i = lines2[i][1], votes_i = lines2[i][2];
        cv::Point pt1, pt2;
        double a = std::cos(theta_i), b = std::sin(theta_i);
        double x0 = a*rho_i, y0 = b*rho_i;
        pt1.x = std::round(x0 + 1000*(-b));
        pt1.y = std::round(y0 + 1000*a);
        pt2.x = std::round(x0 - 1000*(-b));
        pt2.y = std::round(y0 - 1000*a);
        cv::line(dst2, pt1, pt2, cv::Scalar(255), 3, cv::LINE_AA);
    }
    std::cout << "================================" << std::endl;
    
    std::cout << (_is_result_same(lines, lines2) ? "Results are the same." : "Results are not the same.") << std::endl;
    
    cv::imshow("src", src);
    cv::imshow("binary", binary);
    cv::imshow("dst", dst);
    cv::imshow("dst2", dst2);
    cv::waitKey(0);
    
    return 0;
}
